import numpy as np
import librosa
import onnxruntime as ort
import _pickle as pickle
from numpy.lib.stride_tricks import as_strided
import lzma
from scipy.fft import rfft


class Speaker:
    def __init__(self):
        self.resample_rate = 16000
        self.cnceleb_resnet34_path = "cnceleb_resnet34.onnx"
        cnceleb_resnet34_so = ort.SessionOptions()
        cnceleb_resnet34_so.inter_op_num_threads = 1
        cnceleb_resnet34_so.intra_op_num_threads = 1
        cnceleb_resnet34_so.enable_cpu_mem_arena = True
        cnceleb_resnet34_so.enable_mem_pattern = True
        self.cnceleb_resnet34_session = ort.InferenceSession(
            self.cnceleb_resnet34_path,
            providers=[
                'CUDAExecutionProvider',
                # 'CPUExecutionProvider'
            ], sess_options=cnceleb_resnet34_so)
        self.mel_basis = pickle.load(lzma.open("resources/wepeaker_mel_basis.pkl", "rb"))
        self.window = pickle.load(lzma.open("resources/wepeaker_window.pkl", "rb"))

    def extract_embedding_from_pcm_onnx(self, wav):
        print("输入音频check------------------------------------形状:{},求和:{}".format(wav.shape, wav.sum()))
        embeddings = np.zeros(256)
        for i in range(0, len(wav), 16000 * 3):
            wav_ = wav[i:i + 16000 * 3]
            weight = len(wav_) / 48000
            wav_ = np.pad(wav_, (0, 16000 * 3 - len(wav_)), "constant", constant_values=0)
            wav_ = wav_ * 32768
            print("fbank前输入音频check------------------------------------形状:{},求和:{}".format(
                wav_.shape, wav_.sum()))
            feats = self.fbank_numpy(wav_)
            print("fbank结果check------------------------------------形状:{},求和:{}".format(
                feats.shape, feats.sum()))
            feats = feats - np.mean(feats, 0)
            feats = feats[None] * 1
            print("feats check------------------------------------形状:{},求和:{}".format(
                feats.shape, feats.sum()))
            embeddings = embeddings + self.cnceleb_resnet34_session.run(
                output_names=['embs'], input_feed={'feats': feats})[0][0] * weight
            print("embeddings check------------------------------------形状:{},求和:{}".format(embeddings.shape,
                                                                                               embeddings.sum()))
        return embeddings


    def fbank_numpy(self, waveform_):
        mel_basis_, window_ = self.mel_basis, self.window
        preemphasis_coefficient = 0.97
        use_power = True
        window_shift = 160
        window_size = 400
        padded_window_size = 512
        # --------------------------------------------------------------------------------------------
        num_samples = len(waveform_)
        m = 1 + (num_samples - window_size) // window_shift
        strided_input_ = as_strided(waveform_, shape=(m, window_size),
                                    strides=(waveform_.strides[0] * window_shift, waveform_.strides[0]))
        print("fbank步骤1结果check------------------------------------形状:{},求和:{}".format(strided_input_.shape,
                                                                                              strided_input_.sum()))
        # --------------------------------------------------------------------------------------------
        strided_input_ = strided_input_ - np.mean(strided_input_, axis=1).reshape(-1, 1)
        print("fbank步骤2结果check------------------------------------形状:{},求和:{}".format(strided_input_.shape,
                                                                                              strided_input_.sum()))
        # ---------------------------------------------------------------------------------------------
        offset_strided_input_ = np.pad(strided_input_, ((0, 0), (1, 0)), mode="edge")
        print(
            "fbank步骤3结果check------------------------------------形状:{},求和:{}".format(offset_strided_input_.shape,
                                                                                            offset_strided_input_.sum()))
        # ---------------------------------------------------------------------------------------------
        strided_input_ = strided_input_ - preemphasis_coefficient * offset_strided_input_[:, :-1]
        print("fbank步骤4结果check------------------------------------形状:{},求和:{}".format(strided_input_.shape,
                                                                                              strided_input_.sum()))
        # ---------------------------------------------------------------------------------------------
        strided_input_ = strided_input_ * window_
        print("fbank步骤5结果check------------------------------------形状:{},求和:{}".format(strided_input_.shape,
                                                                                              strided_input_.sum()))
        # ---------------------------------------------------------------------------------------------
        padding_right = padded_window_size - window_size
        strided_input_ = np.pad(strided_input_, ((0, 0), (0, padding_right)), mode="constant", constant_values=0)
        print("fbank步骤6结果check------------------------------------形状:{},求和:{}".format(strided_input_.shape,
                                                                                              strided_input_.sum()))
        # ---------------------------------------------------------------------------------------------
        # spectrum_ = np.abs(scipy.fft.rfft(strided_input_, axis=1))
        spectrum_ = np.abs(rfft(strided_input_, axis=1)).astype(np.float32)
        print("fbank步骤7结果check------------------------------------形状:{},求和:{}".format(spectrum_.shape,
                                                                                              spectrum_.sum()))
        if use_power:
            spectrum_ = np.power(spectrum_, 2)
        mel_energies_ = mel_basis_.dot(spectrum_.T).T
        print("fbank步骤8结果check------------------------------------形状:{},求和:{}".format(mel_energies_.shape,
                                                                                              mel_energies_.sum()))
        mel_energies_ = np.log(np.maximum(mel_energies_, 1.1921e-07))
        return mel_energies_






    @staticmethod
    def to_rmss(wav, hop_size=160, window_size=1024):
        f_nums = len(wav) // hop_size
        wav_ = np.pad(wav, window_size // 2, "reflect")
        new_rmss = np.array([wav_[i * hop_size:i * hop_size + window_size] for i in range(f_nums)])
        new_rmss = np.sqrt(np.mean(new_rmss ** 2, 1))
        return new_rmss




    def run(self, wav):
        rmss = self.to_rmss(wav)
        setence_embeddings = []
        setence_rmss_weight = []
        for i in range(0, len(wav), 16000 * 6):
            wav_piece = wav[i:i + 16000 * 6]
            embedding = self.extract_embedding_from_pcm_onnx(wav_piece)
            setence_embeddings.append(embedding / np.linalg.norm(embedding))
            setence_rmss_weight.append(np.mean(rmss[i // 160:i // 160 + 600]))
        setence_embeddings = np.array(setence_embeddings)
        sim_matrix = setence_embeddings.dot(setence_embeddings.T)
        max_cluster = (sim_matrix > 0.6).sum(1).argmax()
        kernel_embedding = setence_embeddings[sim_matrix[max_cluster] > 0.6].mean(0)
        kernel_embedding = kernel_embedding / np.linalg.norm(kernel_embedding)
        for i in range(5):
            kernel_dis = kernel_embedding.dot(setence_embeddings.T) > 0.6
            print(i, kernel_dis.sum() / len(kernel_dis))
            kernel_embedding = setence_embeddings[kernel_dis].mean(0)
            kernel_embedding = kernel_embedding / np.linalg.norm(kernel_embedding)
        kernel_dis = kernel_embedding.dot(setence_embeddings.T)
        setence_rmss_weight = np.array(setence_rmss_weight)
        purity = setence_rmss_weight[kernel_dis > 0.6].sum() / setence_rmss_weight.sum()
        return purity


if __name__ == '__main__':
    self = Speaker()
    self.cnceleb_resnet34_path = "cnceleb_resnet34.onnx"
    audio_path = "resources/北京欢迎你.WAV"  # 后来
    wav = librosa.load(audio_path, sr=16000)[0]
    embedding = self.extract_embedding_from_pcm_onnx(wav[:16000 * 3])
    purity = self.run(wav)
    print(f"音频文件 {audio_path} 的单人纯净度: {purity}")
